import re
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_selection import SelectKBest, chi2
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.model_selection import train_test_split


# -----------------------
# 1. 自定义特征提取器（邮件结构特征）
# -----------------------
class EmailStructureFeatures(BaseEstimator, TransformerMixin):
    """提取邮件结构特征：如是否包含URL、电话号码、邮箱地址、特殊符号比例等"""

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        features = []
        for text in X:
            # 是否包含URL
            has_url = 1 if re.search(r'http[s]?://|www\.', text) else 0
            # 是否包含电话号码
            has_phone = 1 if re.search(r'1\d{10}|\d{3,4}-\d{7,8}', text) else 0
            # 是否包含邮箱地址
            has_email = 1 if re.search(r'[\w\.-]+@[\w\.-]+', text) else 0
            # 特殊符号比例（如@、$、%等）
            special_chars = len(re.findall(r'[^\w\s]', text)) / (len(text) + 1)  # 避免除零
            # 换行符数量（反映格式复杂度）
            line_breaks = text.count('\n') / (len(text) + 1)

            features.append([has_url, has_phone, has_email, special_chars, line_breaks])
        return np.array(features)


# -----------------------
# 2. 文本向量化与特征管道
# -----------------------
def build_pipeline(use_tfidf=True, select_k=1000):
    """
    构建特征处理管道：
    - 文本向量化（TF-IDF或BOW）
    - 结构特征提取
    - 特征选择（卡方检验/RF重要性）
    """
    # 文本向量化组件
    if use_tfidf:
        text_vectorizer = TfidfVectorizer(
            max_features=2000,
            stop_words='english',
            ngram_range=(1, 2)  # 包含单字和双字短语
        )
    else:
        text_vectorizer = CountVectorizer(
            max_features=2000,
            stop_words='english',
            ngram_range=(1, 2)
        )

    # 特征联合：文本特征 + 结构特征
    feature_union = FeatureUnion([
        ('text_features', text_vectorizer),
        ('structure_features', EmailStructureFeatures())
    ])

    # 特征选择：卡方检验（分类任务常用）
    feature_selector = SelectKBest(chi2, k=select_k)

    # 构建管道
    pipeline = Pipeline([
        ('features', feature_union),
        ('selector', feature_selector)
    ])

    return pipeline


# -----------------------
# 3. 主流程：加载数据并应用管道
# -----------------------
def main(csv_path):
    # 加载预处理后的邮件数据
    df = pd.read_csv(csv_path)
    df = df.dropna(subset=['body', 'label'])  # 过滤无效数据
    X = df['body'].astype(str)
    y = df['label']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    # 初始化管道（使用TF-IDF，选择Top 1000特征）
    pipeline = build_pipeline(use_tfidf=True, select_k=1000)

    # 拟合管道（训练集）
    print("开始特征处理与选择...")
    X_train_processed = pipeline.fit_transform(X_train, y_train)
    X_test_processed = pipeline.transform(X_test)

    print(f"处理后训练集特征维度: {X_train_processed.shape}")
    print(f"处理后测试集特征维度: {X_test_processed.shape}")

    # 可选：查看特征选择结果（以RF重要性为例）
    # 修改点：通过 named_steps 访问 features 步骤，再取 transformer_list
    if hasattr(pipeline.named_steps['features'].transformer_list[0][1], 'get_feature_names_out'):
        # 获取文本特征名称
        text_feature_names = pipeline.named_steps['features'].transformer_list[0][1].get_feature_names_out()
        structure_feature_names = ['has_url', 'has_phone', 'has_email', 'special_char_ratio', 'line_break_ratio']
        all_feature_names = list(text_feature_names) + structure_feature_names

        # 训练RF计算特征重要性
        rf = RandomForestClassifier(n_estimators=100, random_state=42)
        rf.fit(X_train_processed, y_train)

        # 输出Top 20重要特征
        importances = rf.feature_importances_
        indices = np.argsort(importances)[-20:][::-1]
        print("\nTop 20重要特征:")
        for i in indices:
            print(f"{all_feature_names[i]}: {importances[i]:.4f}")

    return X_train_processed, X_test_processed, y_train, y_test, pipeline


# -----------------------
# 执行入口
# -----------------------
if __name__ == '__main__':
    csv_path = 'trec06c_processed.csv'  # 替换为你的CSV路径
    X_train, X_test, y_train, y_test, pipeline = main(csv_path)
    # 后续可直接使用X_train, X_test训练